# -*- coding: utf-8 -*-
"""
Author: chen xi (modified by Qwen)
Date: 2024/12/24 下午2:30
File: mass_tcp_client.py
功能：高并发 TCP 客户端，发送加密数据，支持随机断开重连
"""

import asyncio
import random
import string
from string_enc import iPN_StringEnc


async def connect_and_send_periodically(host, port, interval, message):
    """
    单个连接：周期发送 + 随机断开重连
    """
    while True:  # 外层循环：支持断线重连
        try:
            print(f"🔄 尝试连接 {host}:{port}")
            reader, writer = await asyncio.open_connection(host, port)
            addr = writer.get_extra_info('peername')
            print(f"✅ 成功连接到 {addr}")

            # 确保 message 是 bytes
            if isinstance(message, str):
                send_data = message.encode('utf-8')
            else:
                send_data = message

            # 发送任务
            async def send_data_task():
                while True:
                    try:
                        writer.write(send_data)
                        await writer.drain()
                        print(f"📤 已发送 {len(send_data)} 字节 到 {addr}")
                        await asyncio.sleep(interval)
                    except Exception as e:
                        if not writer.is_closing():
                            print(f"❌ 发送任务中断: {e}")
                        break

            # 启动发送任务
            sender_task = asyncio.create_task(send_data_task())

            # 随机存活时间：30~300 秒
            disconnect_delay = random.uniform(30, 300)
            print(f"⏳ 本连接将在 {disconnect_delay:.1f} 秒后断开")

            try:
                # 等待随机时间后断开（或被发送任务中断）
                await asyncio.wait_for(asyncio.shield(sender_task), timeout=disconnect_delay)
            except asyncio.TimeoutError:
                print(f"⏰ 超时触发，准备断开连接 {addr}")
            except Exception:
                pass  # 发送已失败，直接断开

            # 取消发送任务
            if not sender_task.done():
                sender_task.cancel()
                try:
                    await sender_task
                except asyncio.CancelledError:
                    pass

            # 关闭连接
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()
            print(f"🔌 已断开连接 {addr}")

            # 随机重连延迟：5~30 秒
            reconnect_delay = random.uniform(5, 30)
            print(f"🛌 等待 {reconnect_delay:.1f} 秒后重连...")
            await asyncio.sleep(reconnect_delay)

        except Exception as e:
            print(f"❌ 连接异常: {e}")
            # 避免密集重试
            await asyncio.sleep(random.uniform(5, 15))


async def create_connections(server_id, num_connections, host, port, interval, did_end=None):
    """
    创建多个可随机断开重连的连接
    """
    tasks = []
    if did_end is None or not did_end:
        did_end = ''.join(random.choices(string.ascii_uppercase, k=6))

    for i in range(server_id * num_connections, (server_id + 1) * num_connections):
        # 构造 DID
        did = f"IOTGDD-{i:06d}-{did_end}"
        print(f"🔑 DID: {did}")
        _, encrypted_message = iPN_StringEnc("RSIOTGWakeupKey@", did, 256)

        # ✅ 确保 encrypted_message 是 bytes
        if isinstance(encrypted_message, str):
            try:
                # 假设是 hex 字符串（根据你实际返回值调整）
                encrypted_message = bytes.fromhex(encrypted_message)
            except ValueError:
                try:
                    # 尝试 base64
                    import base64
                    encrypted_message = base64.b64decode(encrypted_message)
                except Exception:
                    # 最后兜底
                    encrypted_message = encrypted_message.encode('utf-8')

        # 启动一个可重连的连接任务
        task = asyncio.create_task(
            connect_and_send_periodically(host, port, interval, encrypted_message)
        )
        tasks.append(task)

        # 防连接风暴
        if i % 10 == 0:
            await asyncio.sleep(0.1)  # 每 10 个连接暂停 100ms

    # 永久运行所有任务（除非手动中断）
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    num_connections = 1      # 每批连接数
    server_id = 0             # 服务器 ID（用于 DID 分区）
    did_end = "AAAAAA"            # 自动生成 6 位后缀
    host = '47.107.56.10'
    port = 50000
    interval = 30             # 心跳间隔（秒）

    print(f"🎯 启动 {num_connections} 个随机断连客户端...")
    print(f"   服务器: {host}:{port}")
    print(f"   心跳: {interval}s | 断连: 30~300s | 重连: 5~30s")
    print("-" * 50)

    try:
        asyncio.run(create_connections(server_id, num_connections, host, port, interval, did_end))
    except KeyboardInterrupt:
        print("\n👋 客户端已停止")